#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from pydantic import BaseModel, Field
from typing import List, Optional, Dict
from fastapi import UploadFile, File

DOC_CUT_OUTPUY_TYPE = ["text", "chunk"]


class FragmentCutRequest(BaseModel):
    query: str = Field(..., example="新能源发展", description="用户问题")
    texts: str = Field(..., example="text1", description="待处理的文本内容")
    chunk_size: Optional[int] = Field(
        default=500,
        example=1024,
        description="设置文档切片大小"
    )
    chunk_overlap: Optional[int] = Field(
        default=50,
        example=50,
        description="切片重叠字数"
    )
    chunk: Optional[bool] = Field(
        default=False,
        example="text",
        description="返回结果类型, text/chunk"
    )
    max_length: Optional[int] = Field(
        default=8192,  # 修正后的单默认值
        example=9192,
        description="设置文档返回内容的最大长度"
    )
    strategy: Optional[str] = Field(
        default="recursive",
        example="recursive",
        description="分片策略，recursive"
    )

    class Config:
        schema_extra = {
            "example": {
                "query": "新能源发展",
                "texts": "这里是要处理的文档内容",
                "max_length": 9192
            }
        }


class FragmentDocExtractLocalRequest(BaseModel):
    query: str = Field(..., example="Hello Word!", description="用户问题")
    files: Optional[List[UploadFile]] = File(
        default=None,
        description="上传的文件（与url参数二选一）",
        max_items=5,
    )
    chunk: Optional[bool] = Field(
        default=False,
        example="text",
        description="是否返回切片列表，默认为False"
    )
    max_length: Optional[int] = Field(
        default=8192,
        example=8192,
        description="设置文档返回内容的最大长度"
    )
    chunk_size: Optional[int] = Field(
        default=500,
        example=1024,
        description="设置文档切片大小"
    )
    chunk_overlap: Optional[int] = Field(
        default=50,
        example=50,
        description="切片重叠字数"
    )
    strategy: Optional[str] = Field(
        default="recursive",
        example="recursive",
        description="分片策略，recursive"
    )

    class Config:
        schema_extra = {
            "example": {
                "query": "Hello Word!",
                "transfer_method": "remote_url",
                "url": "http://example.com/test.docx",
                "max_length": 8192,
                "chunk_size": 500,
                "chunk_overlap": 50,
                "chunk": False
            }
        }


class UrlItem(BaseModel):
    url: str = Field(..., description="文件URL地址")
    extension: str = Field(..., description="文件扩展名(如: docx, pdf等)")
    name: Optional[str] = Field(
        default=None,
        description="文件名，如未指定则使用URL中的文件名"
    )


class FragmentDocExtractUrlRequest(BaseModel):
    query: str = Field(..., example="新能源发展", description="用户问题")
    urls: List[UrlItem] = Field(
        ...,
        description="远程文件URL列表，每个URL需指定文件扩展名",
        example=[{
            "url": "http://example.com/test.docx",
            "extension": "docx"
        }],
        max_items=5
    )
    chunk: Optional[bool] = Field(
        default=False,
        example="text",
        description="是否返回切片列表，默认为False"
    )
    max_length: Optional[int] = Field(
        default=8192,
        example=8192,
        description="设置文档返回内容的最大长度"
    )
    chunk_size: Optional[int] = Field(
        default=500,
        example=1024,
        description="设置文档切片大小"
    )
    chunk_overlap: Optional[int] = Field(
        default=50,
        example=50,
        description="切片重叠字数"
    )
    strategy: Optional[str] = Field(
        default="recursive",
        example="recursive",
        description="分片策略，recursive"
    )

    class Config:
        schema_extra = {
            "example": {
                "query": "Hello Word!",
                "transfer_method": "remote_url",
                "url": "http://example.com/test.docx",
                "max_length": 8192,
                "chunk_size": 500,
                "chunk_overlap": 50,
                "chunk": False
            }
        }
